import os
import numpy as np
from scipy.spatial import KDTree


def filter_and_replace_anomalies(flow, threshold=20, k=1):
    h, w, _ = flow.shape

    # 找到异常值位置
    anomalies = np.abs(flow) > threshold
    anomalies = anomalies[:, :, 0] | anomalies[:, :, 1]

    # 提取正常值的位置和对应的光流值
    normal_positions = np.argwhere(~anomalies)
    normal_values = flow[~anomalies]

    # 检查正常值是否为空
    if normal_positions.size == 0 or normal_values.size == 0:
        raise ValueError("没有找到正常值")

    # 创建KDTree以便快速查找最近邻
    kdtree = KDTree(normal_positions)

    # 替换异常值
    for y in range(h):
        for x in range(w):
            if anomalies[y, x]:
                # 找到最近邻的k个正常值
                _, idx = kdtree.query((y, x), k=k)
                if k == 1:
                    nearest_values = normal_values[idx]
                else:
                    nearest_values = np.mean(normal_values[idx], axis=0)
                flow[y, x] = nearest_values

    return flow


def process_flow_files(input_dir, output_dir, threshold=20, k=1):
    os.makedirs(output_dir, exist_ok=True)
    flow_files = [f for f in os.listdir(input_dir) if f.endswith('.npy')]

    for flow_file in flow_files:
        flow_path = os.path.join(input_dir, flow_file)
        output_flow_path = os.path.join(output_dir, flow_file)

        try:
            flow = np.load(flow_path)
            filtered_flow = filter_and_replace_anomalies(flow, threshold, k)
            np.save(output_flow_path, filtered_flow)
            print(f"Processed {flow_file} and saved to {output_flow_path}")
        except Exception as e:
            print(f"Error processing {flow_file}: {e}")


# 示例用法
input_dir = 'C:\\Users\\crxc\\Pictures\\test\\data\\4096_image_pair\\flow_combine_512'  # 替换为实际输入目录路径
output_dir = 'C:\\Users\\crxc\\Pictures\\test\\data\\4096_image_pair\\flow_combine_512_repair'  # 替换为实际输出目录路径

process_flow_files(input_dir, output_dir, threshold=10, k=1)

print("Batch processing completed.")
